#!/usr/bin/env python
"""
Simple command-line barrier using AMQP.
Both the total number of processes and
the rank of each process must be given as input.

Sample usage (10 hosts):
  host1% amqp_barrier.py localhost 10 0
  host2% amqp_barrier.py localhost 10 1
  ...
  host10% amqp_barrier.py localhost 10 9
"""
__author__ = "Dan Gunter <dkgunter@lbl.gov>"
__rcsid__ = "$Id$"
import os
import signal
import socket
import sys
import time
#
from amqplib import client_0_8 as amqp

EXCHANGE = 'barrier_ex'
QUEUE = 'barrier_q'

# Number of messages that we are still
# waiting for, or None if not yet waiting.
# This is a global for signal handling.
g_count, g_waiting = None, None

# Channel
g_channel = None

# Verbosity flag
g_verbose = False

g_my_id = -1

def usage():
    prog = os.path.basename(sys.argv[0])
    sys.stderr.write("usage: %s [-v] HOST NPROC PID\n"
                     "-v = Verbose\n"
                     "HOST = rabbitmq host\n" 
                     "NPROC = number of processes\n"
                     "PID = my process (0 .. NPROC-1)\n"
                     % prog)

def interrupted(signo, frame):
    """Signal handler.
    Perform clean shutdown.
    """
    verbose("shutting down on signal %d" % signo)
    shutdown()

def shutdown():
    try:
        _shutdown()
    except Exception, err:
        pass
        
def _shutdown():
    """Actions for clean shutdown
    """
    if g_waiting is None:
        sys.exit(0) # stopped before we started
    if g_channel and (g_waiting > 0):
        for i in xrange(g_count):
            q = queue(i)
            verbose("deleting queue %s" % q)
            try:
                g_channel.queue_delete(queue=q, nowait=True)
            except amqp.AMQPException, err:
                verbose("** error deleting queue %s: %s" % (q, err))
        g_channel.close()
        abort(1, "%d nodes waiting" % g_waiting)
    else:
        sys.exit(0) # done, anyways
        
def connect(host):
    """Connect to AMQP server on 'host'.
    """
    conn = amqp.Connection(host=host)
    channel = conn.channel()
    channel.exchange_declare(exchange=EXCHANGE, type='direct',
                             durable=False, auto_delete=True)
    return channel

def hello():
    """Construct and return a hello msg
    """
    body = "hello(%d)" % os.getpid()
    return amqp.Message(body)

def queue(i):
    """Get name of queue number 'i' (0-based).
    """
    return "%s%d" % (QUEUE, i)

def got_one(msg):
    """Action to perform on message receipt.
    """
    global g_waiting
    verbose("got msg: %s" % msg.body)
    g_waiting -= 1
    
def verbose(msg):
    """Print messages if verbose mode is turned on.
    """
    if g_verbose:
        sys.stderr.write("%d: %s\n" % (g_my_id, msg))
        
def abort(code, msg):
    """Abort with given exit code and message.
    """
    sys.stderr.write("** " + msg + "\n")
    sys.stderr.write("Abort.\n")
    sys.exit(code)

def run(my_id):
    """Do the main actions for the barrier.
    """
    global g_waiting
    chan = g_channel
    # declare one queue per consumer
    for i in xrange(g_waiting):
        rkey = q = queue(i)
        chan.queue_declare(queue=q, durable=False, auto_delete=True)
        chan.queue_bind(queue=q, exchange=EXCHANGE, routing_key=rkey)
    # send one message to each queue
    for i in xrange(g_waiting):
        rkey = queue(i)
        verbose("say hello to %s" % rkey)
        chan.basic_publish(hello(), exchange=EXCHANGE, routing_key=rkey)
    # wait for the right number of messages on 'my' queue
    my_q = queue(my_id)
    chan.basic_consume(queue=my_q, callback=got_one, no_ack=True)
    while g_waiting > 0:
        chan.wait()
        
def main():
    """Program entry point.
    """
    global g_waiting, g_count, g_verbose, g_channel, g_my_id
    for sig in signal.SIGINT, signal.SIGHUP, signal.SIGUSR1, signal.SIGUSR2, \
        signal.SIGTERM, signal.SIGPIPE, signal.SIGIO:
        signal.signal(sig, interrupted)
    a, g_verbose = 1, False
    try:
        host = sys.argv[a]
        if host == '-v':
            g_verbose = True
            a += 1
            host = sys.argv[a]
        a += 1
        g_count = g_waiting = int(sys.argv[a])
        a += 1
        id_ = int(sys.argv[a])
    except (IndexError, ValueError), err:
        usage()
        return 1
    if id_ < 0 or id_ > (g_count - 1):
        usage()
        abort(2, "PID not in range 0..NPROC-1")
    # set id globally, mostly for logging
    g_my_id = id_
        
    verbose("start (%d)" % os.getpid())
    while g_channel is None:
        try:
            g_channel = connect(host)
        except Exception, err:
            verbose("cannot connect to '%s': %s" % (host, err))
            verbose("sleep, then reconnect")
            time.sleep(5)
    try:
        run(id_)
    except ValueError, err:
        abort(2, err)
    except (amqp.AMQPException, TypeError, socket.error), err:
        verbose("shutting down on error: %s" % err)
        shutdown()
    g_channel.close()
    verbose("end (%d)" % os.getpid())
    return 0

if __name__ == '__main__':
   sys.exit(main())
